#include "FirstFollowGenerator.h"
#include "ContextFreeGrammar.h"
#include "OutputHelper.h"
#include <algorithm>
#include <iostream>
#include <cassert>

using std::cerr;
using std::endl;

bool FirstFollowGenerator::MergeListKeepUnique(FirstFollowSet& to, const FirstFollowSet& from)
{
    bool hasChanged = false;
    for(auto itr = from.begin(); itr != from.end(); ++itr)
    {
        bool found = false;

        for(auto jtr = to.begin(); jtr != to.end(); ++jtr)
            if(*jtr == *itr)
            {
                found = true;
                break;
            }

        if(!found)
            to.push_back(*itr), hasChanged = true;
    }
    return hasChanged;
}


FirstFollowGenerator::FirstFollowGenerator(ContextFreeGrammar * _parent):
    parent(_parent)
{
    //ctor
}

FirstFollowGenerator::~FirstFollowGenerator()
{
    //dtor
}

void FirstFollowGenerator::Generate(const std::string & startSymbol)
{
    this->GenerateFirst();
    this->PrintFirst(cerr);
    this->GenerateFollow(startSymbol);
    this->PrintFollow(cerr);
    this->GeneratePredTable();
    this->PrintPredTable(cerr);
}

void FirstFollowGenerator::GenerateFirst()
{
    FirstFollowMap ret;

    // Fill ret with empty sets
    for(auto itr = parent->nonterminal.begin(); itr != parent->nonterminal.end(); ++itr)
        ret[*itr] = FirstFollowSet();

    bool hasChanged;
    do
    {
        //cerr << "-Iterating" << endl ;
        hasChanged = false;

        for(auto itr = parent->production.begin(); itr != parent->production.end(); ++itr)
        {

            FirstFollowSet newSet;
            bool allEpsilon = true;
            //cerr << "Processing production:" << itr->first << "->"  ;
            //parent->PrintProductionRight(itr->second);
            //cerr << endl ;

            if(itr->second.size() == 1 && *(itr->second.begin()) == parent->GetEpsilon())
            {
                newSet.push_back(parent->GetEpsilon());
            }
            else
            {
                for(auto ktr = itr->second.begin(); ktr != itr->second.end() && allEpsilon; ++ktr)
                {
                    // If ktr is terminal, then FIRST(ktr) = ktr
                    if(!parent->isNonterminal(**ktr))
                    {
                        newSet.push_back(*ktr);
                        allEpsilon = false;
                    }
                    // If ktr is nonterminal, then for each of FIRST(Y_K)
                    else
                    {
                        bool hasEpsilon = false;
                        for(auto jtr = ret[**ktr].begin(); jtr != ret[**ktr].end(); ++jtr)
                        {
                            newSet.push_back(*jtr);
                            // Determine if epsilon is in FIRST(Y_k)
                            if(*jtr == parent->GetEpsilon())
                                hasEpsilon = true;
                        }
                        if(!hasEpsilon)
                            allEpsilon = false;
                    }
                }
            }

            for(auto jtr = newSet.begin(); jtr != newSet.end(); ++jtr)
            {
                bool found = false;
                for(auto ktr = ret[itr->first].begin(); ktr != ret[itr->first].end(); ++ktr)
                    if(*ktr == *jtr)
                    {
                        found = true;
                        break;
                    }

                if(!found)
                {
                    //cerr << "Inserting " << parent->SafePrint(*jtr) << endl ;
                    ret[itr->first].push_back(*jtr), hasChanged = true;
                }

            }
        }
    }
    while(hasChanged);

    this->first = ret;
}

bool FirstFollowGenerator::CheckEpsilonInFirst(const std::string& str) const
{
    bool hasEpsilon = false;
    if(!parent->isNonterminal(str))
        return false;
    for(auto xitr = first.at(str).begin(); xitr != first.at(str).end(); ++xitr)
        if(parent->GetEpsilon() == *xitr)
        {
            hasEpsilon = true;
            break;
        }
    return hasEpsilon;
}


void FirstFollowGenerator::CollectFirst(FirstFollowSet& oldSet, const std::string& toCollect) const
{
    FirstFollowSet newSet;
    if(!parent->isNonterminal(toCollect))
        newSet.push_back(parent->terminal.find(toCollect));
    else
    {
        const auto & newFirst = this->first.find(toCollect)->second;
        assert(this->first.find(toCollect)->first == toCollect);
        std::for_each(newFirst.begin(), newFirst.end(),
                      [&newSet](const StringSet::const_iterator & toInsert)
        {
            newSet.push_back(toInsert);
        });
    }

    for(auto itr = newSet.begin(); itr != newSet.end(); ++itr)
    {
        bool found = false;
        for(auto jtr = oldSet.begin(); jtr != oldSet.end(); ++jtr)
            if(*itr == *jtr)
            {
                found = true;
                break;
            }
        if(!found)
            oldSet.push_back(*itr);
    }
}


void FirstFollowGenerator::GenerateFollow(const std::string & startSymbol)
{
    //cerr << "Generating follow" << endl ;
    FirstFollowMap ret;
    // Fill ret with empty sets
    for(auto itr = parent->nonterminal.begin(); itr != parent->nonterminal.end(); ++itr)
        ret[*itr] = FirstFollowSet();
    // Put $ into FOLLOW(start symbol)
    if(startSymbol.length() >= 1)
    {
        if(!parent->isNonterminal(startSymbol))
            throw std::runtime_error("Undefined start symbol");
        ret[startSymbol].push_back(parent->GetEndOfFile());
    }
    else
        ret[*parent->startSymbol].push_back(parent->GetEndOfFile());

    bool hasChanged = false;
    do
    {
        //cerr << "-Iterating" << endl ;
        hasChanged = false;
        for(auto pitr = parent->production.begin(); pitr != parent->production.end(); ++pitr)
        {
            //cerr << "Processing production:" << pitr->first << "->"  ;
            //parent->PrintProductionRight(pitr->second);
            //cerr << endl ;


            const auto & production = *pitr;

            if(production.second.size() == 1 && parent->GetEpsilon() == *production.second.begin())
                continue;

            // A -> alpha B beta
            if(production.second.size() >= 2)
            {
                //cerr << "A -> alpha B beta" << endl ;
                for(auto bitr = pitr->second.begin(); bitr != std::prev(pitr->second.end()); ++bitr)
                {
                    if(!parent->isNonterminal(**bitr))
                        continue;

                    FirstFollowSet newSet;
                    //cerr << "B:" << parent->SafePrint(*bitr) << endl ;

                    bool allEpsilon = true;
                    FirstFollowSet betaSet;
                    for(auto betaitr = std::next(bitr); betaitr != pitr->second.end() && allEpsilon; ++betaitr)
                    {
                        CollectFirst(betaSet, **betaitr);
                        allEpsilon = CheckEpsilonInFirst(**betaitr);
                    }
                    if(allEpsilon)
                        betaSet.push_back(parent->GetEpsilon());


                    //cerr << "Collected set has " << betaSet.size() << " elements" << endl;
                    for(auto betaitr = betaSet.begin(); betaitr != betaSet.end(); ++betaitr)
                        if(*betaitr != parent->GetEpsilon())
                            newSet.push_back(*betaitr);
                    /*
                                        for(auto itr = newSet.begin(); itr != newSet.end(); ++itr)
                                        {
                                            bool found = false;
                                            assert(*bitr != parent->GetEpsilon() && *bitr != parent->GetEndOfFile());

                                            for(auto jtr = ret[**bitr].begin(); jtr != ret[**bitr].end(); ++jtr)
                                                if(*jtr == *itr)
                                                {
                                                    found = true;
                                                    break;
                                                }

                                            if(!found)
                                                ret[**bitr].push_back(*itr), hasChanged = true;
                                        }*/
                    bool changed = MergeListKeepUnique(ret[**bitr], newSet);
                    if(changed)
                        hasChanged = true;
                }
            }

            // A -> alpha B
            if(production.second.size() >= 1 &&
                    parent->isNonterminal(**std::prev(production.second.end())))
            {
                FirstFollowSet newSet;
                //cerr << "A -> alpha B" << endl ;
                std::for_each(ret[pitr->first].begin(), ret[pitr->first].end(),
                              [&newSet](StringSet::const_iterator itr)
                {
                    newSet.push_back(itr);
                });

                auto bitr = std::prev(production.second.end());
                /*
                for(auto itr = newSet.begin(); itr != newSet.end(); ++itr)
                {
                    bool found = false;
                    assert(*bitr != parent->GetEpsilon() && *bitr != parent->GetEndOfFile());

                    for(auto jtr = ret[**bitr].begin(); jtr != ret[**bitr].end(); ++jtr)
                        if(*jtr == *itr)
                        {
                            found = true;
                            break;
                        }

                    if(!found)
                        ret[**bitr].push_back(*itr), hasChanged = true;
                }*/
                bool changed = MergeListKeepUnique(ret[**bitr], newSet);
                if(changed)
                    hasChanged = true;
            }

            // A -> alpha B beta & epsilon in FIRST(beta)
            if(production.second.size() >= 2)
            {
                //cerr << "A -> alpha B beta & epsilon" << endl ;
                for(auto bitr = pitr->second.begin(); bitr != std::prev(pitr->second.end()); ++bitr)
                {
                    if(!parent->isNonterminal(**bitr))
                        continue;

                    FirstFollowSet newSet;
                    //cerr << "B:" << parent->SafePrint(*bitr) << endl;

                    bool allEpsilon = true;
                    FirstFollowSet betaSet;
                    for(auto betaitr = std::next(bitr); betaitr != pitr->second.end() && allEpsilon; ++betaitr)
                    {
                        CollectFirst(betaSet, **betaitr);
                        //if(!CheckEpsilonInFirst(**betaitr))
                        allEpsilon = CheckEpsilonInFirst(**betaitr);
                    }
                    if(allEpsilon)
                        betaSet.push_back(parent->GetEpsilon());
                    /*
                                        bool hasEpsilon = false;
                                        for(auto betaitr = betaSet.begin(); betaitr != betaSet.end(); ++betaitr)
                                            if(*betaitr == parent->GetEpsilon())
                                            {
                                                hasEpsilon = true;
                                                break;
                                            }
                    */
                    if(allEpsilon)
                    {
                        std::for_each(ret[pitr->first].begin(), ret[pitr->first].end(),
                                      [&newSet](StringSet::const_iterator itr)
                        {
                            newSet.push_back(itr);
                        });
                    }
                    /*
                    for(auto itr = newSet.begin(); itr != newSet.end(); ++itr)
                    {
                        bool found = false;
                        assert(*bitr != parent->GetEpsilon() && *bitr != parent->GetEndOfFile());

                        for(auto jtr = ret[**bitr].begin(); jtr != ret[**bitr].end(); ++jtr)
                            if(*jtr == *itr)
                            {
                                found = true;
                                break;
                            }

                        if(!found)
                            ret[**bitr].push_back(*itr), hasChanged = true;
                    }
                    */
                    bool changed = MergeListKeepUnique(ret[**bitr], newSet);
                    if(changed)
                        hasChanged = true;
                }
            }
        }
    }
    while(hasChanged);
    follow = ret;
}

void FirstFollowGenerator::GeneratePredTable()
{
    //cerr << "Generating Prediction table" << endl ;
    for(auto pitr = parent->production.begin(); pitr != parent->production.end(); ++pitr)
    {
        //cerr << "Processing production:" << pitr->first << "->"  ;
        //parent->PrintProductionRight(pitr->second);
        //cerr << endl ;

        const ProductionList & pright = pitr->second;

        bool AlphaProduction = (pright.size() == 1 && *pright.begin() == parent->GetEpsilon());


        bool allEpsilon = true;
        if(!AlphaProduction)
        {
            FirstFollowSet firstAlpha;
            for(auto alphaitr = pright.begin(); alphaitr != pright.end() && allEpsilon; ++alphaitr)
            {
                CollectFirst(firstAlpha, **alphaitr);
                allEpsilon = CheckEpsilonInFirst(**alphaitr);
            }

            for(auto alphaitr = firstAlpha.begin(); alphaitr != firstAlpha.end(); ++alphaitr)
            {
                //cerr << "Inserting " << parent->SafePrint(*alphaitr) << endl ;
                if(*alphaitr == parent->GetEndOfFile() || *alphaitr == parent->GetEpsilon())
                    continue;
                if(!parent->isNonterminal(**alphaitr))
                {
                    if(predTable[pitr->first].find(**alphaitr) != predTable[pitr->first].end())
                        throw std::runtime_error("Ambiguous grammar");
                    predTable[pitr->first][**alphaitr] = pitr;
                }

            }
        }



        // there is Epsilon in FIRST(alpha)
        if(allEpsilon)
        {
            const FirstFollowSet & followAlpha = follow[pitr->first];
            for(auto followitr = followAlpha.begin(); followitr != followAlpha.end(); ++followitr)
            {
                //cerr << "Inserting " << parent->SafePrint(*followitr) << endl ;
                if(*followitr == parent->GetEpsilon())
                    continue;
                if(*followitr == parent->GetEndOfFile())
                    predTable[pitr->first]["$"] = pitr;
                else if(!parent->isNonterminal(**followitr))
                {
                    if(predTable[pitr->first].find(**followitr) != predTable[pitr->first].end())
                        throw std::runtime_error("Ambiguous grammar");
                    predTable[pitr->first][**followitr] = pitr;
                }
            }
        }
    }
}


void FirstFollowGenerator::Print(std::ostream& out) const
{
    out << "First:" << std::endl ;
    this->PrintFirst(out);
    out << "Follow:" << std::endl ;
    this->PrintFollow(out);
    out << "Prediction Map:" << std::endl ;
    this->PrintPredTable(out);
}

void FirstFollowGenerator::PrintFirst(std::ostream& out) const
{
    out << "First:" << endl ;
    for(auto itr = first.cbegin(); itr != first.cend(); ++itr)
    {
        out << itr->first << "\t:" ;
        for(auto jtr = itr->second.begin(); jtr != itr->second.end(); ++jtr)
        {
            out << parent->SafePrint(*jtr) << ";" ;
        }
        out << std::endl ;
    }
}

void FirstFollowGenerator::PrintFollow(std::ostream& out) const
{
    out << "Follow:" << endl ;
    for(auto itr = follow.cbegin(); itr != follow.cend(); ++itr)
    {
        out << itr->first << "\t:" ;
        for(auto jtr = itr->second.begin(); jtr != itr->second.end(); ++jtr)
        {
            out << parent->SafePrint(*jtr) << ";" ;
        }
        out << std::endl ;
    }
}

void FirstFollowGenerator::PrintPredTable(std::ostream & out) const
{
    out << "Prediction Table:" << endl ;
    for(auto termitr = predTable.begin(); termitr != predTable.end(); ++termitr)
    {
        out << termitr->first << "\t:" ;
        for(auto ntermitr = termitr->second.begin(); ntermitr != termitr->second.end(); ++ntermitr)
        {
            out << ntermitr->first << ":";
            parent->PrintProductionRight(ntermitr->second->second);
            out << "\t";
        }
        out << std::endl ;
    }
}

std::pair <const FirstFollowMap *, const FirstFollowMap *> FirstFollowGenerator::GetFirstFollow() const
{
    return std::make_pair(&first, &follow);
}


